'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: base_agent.py
* @Time: 2025/10/9
* @All Rights Reserve By Brtc
'''
import uuid
from abc import ABC, abstractmethod
from threading import Thread
from typing import Optional, Any
from langchain_core.language_models import BaseLanguageModel
from langchain_core.load import Serializable
from langchain_core.runnables import Runnable, RunnableConfig
from langgraph.graph.graph import CompiledGraph
from pydantic.v1 import PrivateAttr
from typing_extensions import Generator
from internal.core.agent.agents.agent_queue_manager import AgentQueueManager
from internal.core.agent.entities.agent_entity import AgentConfig, AgentState
from internal.core.agent.entities.queue_entity import AgentResult
from internal.exception.exception import FailException


class BaseAgent(Serializable, Runnable):
    """llmops项目的基础Agent类"""
    llm:BaseLanguageModel
    agent_config:AgentConfig
    _agent : CompiledGraph = PrivateAttr(None)
    _agent_queue_manager: AgentQueueManager = PrivateAttr(None)

    class Config:
        #字段允许接受任何类类型， 且不需要校验
        arbitrary_types_allowed = True

    def __init__(self,
                 llm:BaseLanguageModel,
                 agent_config:AgentConfig,
                 *args,
                 **kwargs):
        """构造函数, 初始化智能体图构造函数"""
        super().__init__(*args, llm=llm, agent_config=agent_config,**kwargs)
        self.agent_config = agent_config
        self._agent = self._build_agent()
        self._agent_queue_manager = AgentQueueManager(
            user_id=agent_config.user_id,
            invoke_from=agent_config.invoke_from
        )
    @abstractmethod
    def _build_agent(self)->CompiledGraph:
        """构建智能体函数， 这个需要子类去实现"""
        pass

    def invoke(self, input:AgentState, config:Optional[RunnableConfig]=None)->AgentResult:
        """块内容生成函数， 一次性生成所有的答案"""
        pass


    def stream(self,
               input:AgentState,
               config:Optional[RunnableConfig] = None,
               **kwargs:Optional[Any])->Generator:
        """流式输出， 每个node 节点或者LLM每生成一个token则会返回相应的内容"""
        #1、检测子类是否已经构建智能体， 如没有构建则抛出错误
        if not self._agent:
            raise FailException("智能体未创建成功， 请创建后再试！！")
        #2、构建对应的任务id 以及初始化
        input["task_id"] = input.get("task_id", uuid.uuid4())
        input["history"] = input.get("history", [])
        input["iteration_count"] = input.get("iteration_count", 0)
        print("================================1=============================")
        print(input)
        #3、创建子线程并执行
        thread = Thread(
            target=self._agent.invoke,
            args=(input,)
        )
        thread.start()
        #4、调用队列管理器监听数据并返回
        yield from self._agent_queue_manager.listen(input["task_id"])

    @property
    def agent_queue_manager(self)->AgentQueueManager:
        return self._agent_queue_manager